import sentence_transformers
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from chatglm_loader import ChatGLM


embedding_model_dict = {\
    "text2vec": "C:\\Work\\llm\\text2vec-large-chinese"
}

EMBEDDING_MODEL = "text2vec"
# 初始化 hugginFace 的 embeddings 对象
embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[EMBEDDING_MODEL])
embeddings.client = sentence_transformers.SentenceTransformer(
    embeddings.model_name, device='cpu')

db = Chroma(persist_directory="./chroma/casting2", embedding_function=embeddings)


# 选择模型
model = ChatGLM()
model.load_model("C:\\Work\\llm\\ChatGLM2-6B\\THUDM\\chatglm2-6b")

retriever = db.as_retriever()
qa = RetrievalQA.from_chain_type(llm=model, chain_type="stuff", retriever=retriever)

# 进行问答
query = "请用中文回答：东芝压铸机故障部位压铸机大杠抽插不动作的故障原因"
print(qa.run(query))

query = "请用中文回答：东芝压铸机故障部位压铸机安全门报警，无动作的故障原因"
print(qa.run(query))


query = "请用中文回答：东芝压铸机故障部位压铸机模具抽芯无动作的故障原因"
print(qa.run(query))


query = "请用中文回答：东芝压铸机故障部位压铸机机床压力异常，不增压的故障原因"
print(qa.run(query))


query = "请用中文回答：东芝压铸机故障部位压铸机压射前进动作缓慢的故障原因"
print(qa.run(query))


